import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from std_msgs.msg import Int32MultiArray 
import cv2
import numpy as np
from cv_bridge import CvBridge
from collections import deque

class ImageRepublisher(Node):
    def __init__(self):
        super().__init__('image_republisher')
        self.subscription = self.create_subscription(
            Image,
            '/image',
            self.listener_callback,
            10)
        self.xy_publisher = self.create_publisher(Int32MultiArray, '/xy_coordinates', 10)  # 创建新的发布者，用于发布x和y坐标
        self.zero_count = 0
        self.last_x0 = 320
        self.last_y0 = 240
        self.points_queue = deque(maxlen=1) 
        self.bridge = CvBridge()  # 用于将ROS图像消息转换为OpenCV格式

    def listener_callback(self, msg):
        np_arr = np.frombuffer(msg.data, np.uint8)
        cv_image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
        x, y= self.process_frame(cv_image)

        # 创建 Int32MultiArray 消息，用于发布x和y坐标
        xy_msg = Int32MultiArray()
        xy_msg.data = [x, y]

        # 发布x和y坐标
        self.xy_publisher.publish(xy_msg)
        if x!=320:
            print(x, y)
    
    def process_frame(self,frame):

        # 默认值，防止未定义的情况
        x0, y0 = self.last_x0, self.last_y0

        # 定义存储点的数组，最多存储5个点的(x, y)坐标
        #points_arr = np.zeros((5, 2))

        # 定义感兴趣区域（ROI）
        ROI = frame[180:480, 0:640]
        gray = cv2.cvtColor(ROI, cv2.COLOR_BGR2GRAY)
        # 使用Canny边缘检测
        edges = cv2.Canny(ROI, 100, 250)
        #_, otsu_thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        # 查找边缘轮廓
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        # 创建两个空白图像用于填充
        ROI_1 = np.zeros_like(ROI)
        ROI_2 = np.zeros_like(ROI)

        if contours:
            for cnt in contours:
                cv2.fillPoly(ROI_1, [cnt], (0, 0, 255))
                #cv2.drawContours(ROI_1, [cnt], -1, (0, 255, 255), 2)
        #cv2.imshow('ROI_1', ROI_1)
        # 转换为HSV颜色空间
        HSV = cv2.cvtColor(ROI, cv2.COLOR_BGR2HSV)

        # 定义颜色阈值
        lower_threshold = np.array([0, 0, 0])
        upper_threshold = np.array([80, 100, 100])
        # 创建颜色掩膜
        mask = cv2.inRange(HSV, lower_threshold, upper_threshold)
        #inverse_mask = cv2.bitwise_not(mask)
        kernel1 = np.ones((5, 5), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel1, iterations=1)
        ROI_2[mask == 255] = [0, 255, 0]
        #cv2.imshow('ROI_2', ROI_2)
        #cv2.imshow('inverse_mask', inverse_mask)
        # 创建红色掩膜和绿色掩膜
        red_mask = cv2.inRange(ROI_1, np.array([0, 0, 255]), np.array([0, 0, 255]))
        green_mask = cv2.inRange(ROI_2, np.array([0, 255, 0]), np.array([0, 255, 0]))

        # 逻辑AND操作找到交集
        intersection_mask = cv2.bitwise_and(red_mask, green_mask)
        intersection_mask = cv2.morphologyEx(intersection_mask, cv2.MORPH_OPEN, kernel1, iterations=1)
        #cv2.imshow('intersection_mask', intersection_mask)
        # 查找重叠掩膜中的轮廓
        contours, _ = cv2.findContours(intersection_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)

        count = 0
        points = []
        if contours:
            for contour in contours:
                Area = cv2.contourArea(contour)
                if Area >= 100:
                    M = cv2.moments(contour)
                    if M["m00"] != 0:
                        cX = int(M["m10"] / M["m00"])
                        cY = int(M["m01"] / M["m00"]) + 180
                    else:
                        cX, cY = 0, 180
                    points.append((cX, cY))  # 添加有效轨迹块点
                    count += 1

        if count > 0:
            # 获取这一帧的y最小的点并添加到队列中
            min_y_point = min(points, key=lambda p: p[1])
            self.points_queue.append(min_y_point)
            self.zero_count = 0  # 重置丢线计数
        else:
            self.zero_count += 1
            if self.zero_count >= 10 and self.points_queue:
                # 当连续丢线，使用上一次有效点的对称值
                x0, y0 = self.last_x0, self.last_y0
                self.zero_count = 0
            else:
                x0, y0 = self.last_x0, self.last_y0

        if len(self.points_queue) == 1:
            # 在三帧中找到y最小的点
            x0, y0 = min(self.points_queue, key=lambda p: p[1])
            '''
            获取三帧平均值
            avg_x0 = np.mean([p[0] for p in points_queue])
            avg_y0 = np.mean([p[1] for p in points_queue])
            x0, y0 = int(avg_x0), int(avg_y0)
            '''
            self.last_x0, self.last_y0 = x0, y0  # 更新最后有效点
            self.points_queue.popleft()    #移除最早的一帧

        # 绘制结果点
        #cv2.circle(frame, (x0, y0), 5, (0, 255, 0), 3)
        if x0 > 540:
            x0 = 540
        if x0 < 100:
            x0 = 100
        return x0,y0

def main(args=None):
    rclpy.init(args=args)
    xy = ImageRepublisher()
    rclpy.spin(xy)
    rclpy.shutdown()

if __name__ == '__main__':
    main()